﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using XUnity.RuntimeHooker.Core;
using XUnity.RuntimeHooker.Properties;

namespace XUnity.RuntimeHooker
{
   public static class RuntimeMethodPatcher
   {
      private static readonly object Sync = new object();
      private static readonly Dictionary<MethodBase, Type> HookedMethodToTrampolineInitializer = new Dictionary<MethodBase, Type>();

      public static void Patch( MethodBase originalMethod, HookMethod prefix, HookMethod postfix )
      {
         // WARNING: This code is in no way guaranteed to work, as it makes some wild assumptions
         //          about how the underlying machine code is being generated during JIT

         // If the code does not work it will likely crash the application immediately with a cryptic or no error message.

         // This hooking implementation works like this:
         // 1. Create a JMP at the start of the hooked method that goes to a method of an identical signature (trampoline function)
         // 2. When the function is entered, rewrite back the old code that was overwritten in step 1
         // 3. Call prefixes, then call original method, then call postfixes
         // 4. Write back in the JMP that was created in step 1

         // This is all very simple. The primary issue this library solves is how to do this through a 'nice'
         // API surface, similar to that of Harmony so that it can be used as a fallback mechanism.

         // The difficult issue is how to navigate to a method of identical signature to the original one
         // if you cannot emit any IL code (netstandard2.0 limitation) because it is difficult to
         // generate a method with the required signature at compile time, that includes types that may
         // be inaccessible.

         // This library pre-defines a set of generic classes that each has an 'Override' method, that will
         // be used as the trampoline function. Each version of the 'Override' method makes different assumptions
         // about the method signature and uses generic parameters defined on the class in order to do so.

         // The weakness of this approach is that any arguments (generated by the JITTER) in relation to generics
         // will *not* be passed to the override method, as it would usually expect. (I think that with the legacy
         // .NET JITTER in 64-bit, these are often passed as an id in the RCX register of the CPU, but this is likely
         // different for Mono and other architectures). This means that if any code is generated in the Override
         // method that requires this information that the application will immediately fail.

         // This has an especially big impact on the JITTER used in Mono, where accessing any generic arguments
         // in the override method will immediately cause application failure. This imposes a big limitation
         // on this approach in that the current implementation is only able to hook methods that return reference
         // types (or void) because it is unable to use a generic argument (TReturn) to determine the size of the variable
         // to send back to the caller (Should it use the stack? Should it use a CPU register? It does not know!).

         // Another limitation of this in Mono is that you cannot use this approach to access static fields on
         // a per-generic class basis (as that would mean using the generic arguments). To overcome this issue
         // a new assembly that maintains state for each hooked method is generated at runtime (by renaming an
         // existing assembly and loading it into memory). This is essentially what this function does.

         // A better approach than this would be to write the trampoline function in machine code. But that
         // would take a whole lot more effort than this. And require knowledge about how each architecture
         // and platform emits machine code.

         // Attempting to create a StackTrace will also break this implementation.

         lock( Sync )
         {
            var flags = BindingFlags.Static | BindingFlags.NonPublic | BindingFlags.Public;

            // check if we have already hooked the specified method
            if( !HookedMethodToTrampolineInitializer.TryGetValue( originalMethod, out var trampolineInitializer ) )
            {
               // obtain data for the assembly that contains the hooking logic/trampolines
               var assemblySize = Resources.b979b301af8b4ef48f224de6dcf2ddf6.Length;

               // copy the data into a new array
               var assemblyData = new byte[ assemblySize ];
               Buffer.BlockCopy( Resources.b979b301af8b4ef48f224de6dcf2ddf6, 0, assemblyData, 0, assemblySize );

               // rename the assembly, so it can be loaded multiple times (in Mono)
               var newName = Guid.NewGuid().ToString( "N" );
               FindAndReplace( "b979b301af8b4ef48f224de6dcf2ddf6", newName, assemblyData );
               var assembly = Assembly.Load( assemblyData );

               // obtain the trampoline initializer type from the dynamically named assembly
               trampolineInitializer = assembly.GetType( "XUnity.RuntimeHooker.Trampolines.TrampolineInitializer" );

               // call the setup hook method with the originalMethod to hook
               trampolineInitializer.GetMethod( "SetupHook", flags ).Invoke( null, new object[] { originalMethod } );

               // add tbe method to our hooked methods dictionary, so we do not hook it again
               HookedMethodToTrampolineInitializer.Add( originalMethod, trampolineInitializer );
            }

            // after hooking the method, lets apply our prefixes/postfixes
            if( prefix != null )
            {
               trampolineInitializer.GetMethod( "AddPrefix", flags ).Invoke( null, new object[] { prefix } );
            }

            if( postfix != null )
            {
               trampolineInitializer.GetMethod( "AddPostfix", flags ).Invoke( null, new object[] { postfix } );
            }
         }
      }

      private static void FindAndReplace( string ansiStringToSearchFor, string ansiStringToReplaceWith, byte[] data )
      {
         if( ansiStringToReplaceWith.Length != ansiStringToSearchFor.Length )
            throw new ArgumentException( "The string to search for must be same size as string to replace." );

         int pos = 0;
         char c;
         byte b;
         int startIdx = 0;
         int searchLen = ansiStringToReplaceWith.Length;

         for( int i = 0 ; i < data.Length ; i++ )
         {
            b = data[ i ];
            c = (char)b;

            if( pos == searchLen )
            {
               for( int j = 0 ; j < searchLen ; j++ )
               {
                  var r = ansiStringToReplaceWith[ j ];
                  data[ j + startIdx ] = (byte)r;
               }

               pos = 0;
            }
            else if( c == ansiStringToSearchFor[ pos ] )
            {
               if( pos == 0 )
               {
                  startIdx = i;
               }

               pos++;
            }
            else
            {
               pos = 0;
            }
         }
      }
   }
}
